package com.pig4cloud.pigx.knowledge.support.rule;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.pig4cloud.pigx.knowledge.config.properties.AiKnowledgeProperties;
import com.pig4cloud.pigx.knowledge.config.properties.RiskControlProperties;
import com.pig4cloud.pigx.knowledge.dto.ChatMessageContext;
import com.pig4cloud.pigx.knowledge.entity.AiBillEntity;
import com.pig4cloud.pigx.knowledge.service.AiBillService;
import com.yomahub.liteflow.core.NodeComponent;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.Nullable;
import org.springframework.stereotype.Component;

/**
 * 基于IP 的风控规则
 *
 * @author lengleng
 * @date 2024/3/26
 */
@Slf4j
@Component("ip")
@RequiredArgsConstructor
public class IPRiskControlRule extends NodeComponent {

	private final AiKnowledgeProperties aiKnowledgeProperties;

	private final AiBillService aiBillService;

	@Override
	public void process() {
		ChatMessageContext messageContext = this.getContextBean(ChatMessageContext.class);

		String risked = riskControl(messageContext.getIp());
		if (StrUtil.isNotBlank(risked)) {
			messageContext.setErrorMessage(risked);
			// 设置结束
			setIsEnd(true);
		}
	}

	/**
	 * 风控检查
	 * @return 风控提示
	 */
	@Nullable
	private String riskControl(String ip) {
		RiskControlProperties riskControl = aiKnowledgeProperties.getRiskControl();
		if (!riskControl.isEnabled()) {
			return null;
		}

		// 根据IP查询今天的调用次数
		long count = aiBillService.count(Wrappers.<AiBillEntity>lambdaQuery()
			.eq(StrUtil.isNotBlank(ip), AiBillEntity::getIp, ip)
			.ge(AiBillEntity::getCreateTime, DateUtil.yesterday()));

		log.info("ip:{},count:{}", ip, count);
		if (count > riskControl.getTimes()) {
			return riskControl.getErrorMsg();
		}

		return null;
	}

}
